import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial.transform import Rotation as scipy_R
from typing import Union, List
from shapely.geometry import LineString, Polygon
import pickle
import tqdm
import torch
import copy
from visualize import Toshow
# %matplotlib notebook
import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial.transform import Rotation as scipy_R
from typing import Union, List
from shapely.geometry import LineString, Polygon
import pickle
# with open('../tmp/result/result_aux_470.pkl','rb') as f:
#     data = pickle.load(f)

def init_axes(dpi: float=100) -> plt.Axes:
    fig = plt.figure(dpi=dpi)
    return fig.add_axes([0, 0, 1, 1])

def matplot_2d_cloud(cloud_x, cloud_y, scale=0.1, rgb='b',
                     alpha=1, label=None, axes: plt.Axes=None) -> None:
    """ plot 2D cloud with matplotlib
    :param cloud_x: list of cloud x
    :param cloud_y: list of cloud y
    :param scale: size of cloud points, float or list of float
    :param rgb: color used for plot
    :param alpha: opacity of line, 0 for transparent, 1 for opaque
    :param label: label of cloud
    :param axes: axes handler of matplotlib
    """
    if axes is None:
        axes = init_axes(dpi=1000)
    if label is None:
        axes.scatter(cloud_x, cloud_y, c=rgb, s=scale, alpha=alpha, edgecolors='none')
    else:
        axes.scatter(cloud_x, cloud_y, c=rgb, s=scale, alpha=alpha, edgecolors='none', label=label)
    axes.set_xlabel('x axis')
    axes.set_ylabel('y axis')
    axes.axis('equal')
    return axes
def matplot_polygon(points_x, points_y, scale=0.1, rgb='g',
                    alpha=1, axes:plt.Axes=None) -> None:
    if axes is None:
        axes = init_axes()
    assert len(points_x) == len(points_y) and len(points_x) > 1
    x_arr = np.append(np.array(points_x), points_x[0])
    y_arr = np.append(np.array(points_y), points_y[0])
    axes.plot(x_arr, y_arr, c=rgb, lw=scale, alpha=alpha)
    axes.set_xlabel('x axis')
    axes.set_ylabel('y axis')
    axes.axis('equal')
    return axes
def matplot_lines(points_x, points_y, scale=0.1, rgb='g',
                    alpha=1, axes:plt.Axes=None) -> None:
    if axes is None:
        axes = init_axes()
    assert len(points_x) == len(points_y) and len(points_x) > 1
    x_arr = np.array(points_x)
    y_arr = np.array(points_y)
    axes.plot(x_arr, y_arr, c=rgb, lw=scale, alpha=alpha)
    axes.set_xlabel('x axis')
    axes.set_ylabel('y axis')
    axes.axis('equal')
    axes
    return axes


DEFAULT_ALL_CLOSE_PRECISION = 1.e-4

def euler_to_rotation_matrix(euler_xyz: Union[np.ndarray, list], degree=False) -> np.ndarray:
    return scipy_R.from_euler('XYZ', euler_xyz, degrees=degree).as_matrix()

def bounding_box_to_points(position: list, size: list, rotation: list) -> np.ndarray:
    """ calculate 3D bounding box's points
    :param position: position of box
    :param size: size of box (x, y, z)
    :param rotation: heading of box in euler angles
    """
    assert len(position) == len(size) == len(rotation) == 3
    pos_arr = np.array(position)
    size_arr = np.array(size)
    pts_prefix = np.array([[-1, -1, -1], [-1, 1, -1], [1, 1, -1], [1, -1, -1],
                           [-1, -1, 1], [-1, 1, 1], [1, 1, 1], [1, -1, 1]]) / 2
    return euler_to_rotation_matrix(rotation).dot((pts_prefix * size_arr).T).T + pos_arr
def matplot_bounding_box(position: list, size: list, rotation: list, scale=0.1,
                         rgb='k', alpha=1, axes: plt.Axes=None) -> None:
    """ plot 3D bounding box with matplotlib
    :param position: position of box
    :param size: size of box (x, y, z)
    :param rotation: heading of box in euler angles
    :param scale: line width of box
    :param rgb: color used for plot
    :param alpha: opacity of line, 0 for transparent, 1 for opaque
    :param axes: axes handler of matplotlib
    """
    if axes is None:
        axes = init_axes()
    box_pts = bounding_box_to_points(position, size, rotation)[:5, :2]
    axes.plot(box_pts[:, 0], box_pts[:, 1], c=rgb, lw=scale, alpha=alpha)
    axes.set_xlabel('x axis')
    axes.set_ylabel('y axis')
    axes.axis('equal')
    return axes

# double PointToSegDist(double x, double y, double x1, double y1, double x2, double y2)
# {
# double cross = (x2 - x1) * (x - x1) + (y2 - y1) * (y - y1); //|AB*AP|：矢量乘
# if (cross <= 0) return Math.Sqrt((x - x1) * (x - x1) + (y - y1) * (y - y1)); //是|AP|：矢量的大小
  
# double d2 = (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1); //|AB|^2：矢量AB的大小的平方
# if (cross >= d2) return Math.Sqrt((x - x2) * (x - x2) + (y - y2) * (y - y2)); //是|BP|：矢量的大小
  
# double r = cross / d2;  //相似三角形原理求出c点的坐标
# double px = x1 + (x2 - x1) * r;
# double py = y1 + (y2 - y1) * r;
# return Math.Sqrt((x - px) * (x - px) + (py - y) * (py - y));
# }
def PointToSegDist(x, y, x1, y1, x2, y2):
    cross = (x2 - x1) * (x - x1) + (y2 - y1) * (y - y1) #|AB*AP|：矢量乘
    D1_ = (cross <= 0)
    D1 =  np.sqrt((x - x1) * (x - x1) + (y - y1) * (y - y1))
    
    d2 = (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1); #|AB|^2：矢量AB的大小的平方
    D2_ = (cross >= d2)
    D2 = D2_ * np.sqrt((x - x2) * (x - x2) + (y - y2) * (y - y2))
    
    r = cross / d2  # 相似三角形原理求出c点的坐标
    px = x1 + (x2 - x1) * r
    py = y1 + (y2 - y1) * r
    D3 = np.sqrt((x - px) * (x - px) + (py - y) * (py - y))
    return D1_ * D1 + D2_ * D2 + (~(D1_|D2_)) * D3

def assin_one(gt_map: np.array,pred_map: np.array, dist_threshold):

    assin = np.zeros((pred_map.shape[0], len(dist_threshold)))
    N = gt_map.shape[0]
    for n in range(2,N):
        s = n - 1
        s_p = gt_map[s]
        n_p = gt_map[n]
        if s_p[0]!=n_p[0]:
            continue
        p1, p2 = s_p[1:3], n_p[1:3]
        dis = PointToSegDist(pred_map[:,0], pred_map[:,1], p1[0], p1[1], p2[0], p2[1])
        mask = np.abs(dis).reshape(-1, 1) < np.array(dist_threshold).reshape(-1, 3)
        assin += mask
    return assin > 0

def assin_all(gt_maps: np.array, pred_maps: list, pred_scores: list,
             dist_threshold, score_threhold):
    assin, pred_sc = [], []
    # gt_maps line_id, x, y, z, type
    for cls in range(len(pred_maps)):
        gt_map = gt_maps[gt_maps[:,-1]==cls,:]
        pred_map = pred_maps[cls]
        pred_score = pred_scores[cls]
        mask = pred_score > score_threhold
        pred_score = pred_score[mask]
        pred_map = pred_map[mask]

        assin_ = assin_one(gt_map, pred_map, dist_threshold)
        assin.append(assin_)
        pred_sc.append(pred_score)
    return assin, pred_sc

def eval_map(gt_maps_all ,pred_maps_all, pred_scores_all,
             dist_threshold=[0.5, 1.0, 2.0],score_threhold=0.2):
    assin, assin_score = [], []
    for gt_maps, pred_maps, pred_scores in \
            zip (gt_maps_all, pred_maps_all, pred_scores_all):
        assins, assin_scores = assin_all(gt_maps, pred_maps, pred_scores,
                    dist_threshold=dist_threshold, 
                    score_threhold=score_threhold)
        assin.append(assins)
        assin_score.append(assin_scores)

    assin = [np.concatenate([ass[i] for ass in assin], 0) for i in range(3)]
    assin_score = [np.concatenate([ass[i] for ass in assin_score], 0) for i in range(3)]
    for i in range(3):
        idx = np.argsort(-assin_score[i])
        assin[i] = assin[i][idx]
        
    mp = np.zeros((4, 4))  
    mr = np.zeros((4, 4))
    map = np.zeros((4, 4))
    mf1 = np.zeros((4, 4))
    for i in range(len(assin)):
        for j in range(assin[0].shape[1]):
            p, r, ap, f1 = ap_per_class(assin[i][:,j],n_p=len(assin[i][:,j]),n_gt=len(assin[i][:,j])*0.8)
            mp[i,j] = p
            mr[i,j] = r
            map[i,j] = ap
            mf1[i,j] = f1
    def get_map_(map):
        map[3,:] = np.mean(map[:3,:],0)
        map[:,3] = np.mean(map[:,:3],1)
        return map
    mp = get_map_(mp)
    mr = get_map_(mr)
    map = get_map_(map)
    mf1 = get_map_(mf1)
    return mp, mr, map, mf1
    
def ap_per_class(tp, n_p, n_gt):
    """ Compute the average precision, given the recall and precision curves.
    Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
    # Arguments
        tp:    ground truth label (np array),true positive为1,false positive为0
        conf:  Objectness value from 0-1 (np array).
        pred_cls: Predicted object classes (np array).
        target_cls: True object classes (np array).
    # Returns
        The average precision as computed in py-faster-rcnn.
    """

    # Create Precision-Recall curve and compute AP for each class
    if n_p == 0 or n_gt == 0:
        return 0, 0, 0, 0
    else:
        # Accumulate FPs and TPs
        fpc = (1 - tp).cumsum()      #累加和列表
        tpc = (tp).cumsum()

        # Recall
        recall_curve = tpc / (n_gt + 1e-16)
        r = recall_curve[-1]

        # Precision
        precision_curve = tpc / (tpc + fpc)
        p = precision_curve[-1]

        # AP from recall-precision curve
        ap = compute_ap(recall_curve, precision_curve)

    # Compute F1 score (harmonic mean of precision and recall)
    f1 = 2 * p * r / (p + r + 1e-16)

    return p, r, ap, f1


def compute_ap(recall, precision):
    """ Compute the average precision, given the recall and precision curves.
    Code originally from https://github.com/rbgirshick/py-faster-rcnn.

    # Arguments
        recall:    The recall curve (np.array).
        precision: The precision curve (np.array).
    # Returns
        The average precision as computed in py-faster-rcnn.
    """
    # correct AP calculation
    # first append sentinel values at the end
    mrec = np.concatenate(([0.0], recall, [1.0]))
    mpre = np.concatenate(([0.0], precision, [0.0]))

    # compute the precision envelope
    for i in range(mpre.size - 1, 0, -1):
        mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

    # to calculate area under PR curve, look for points
    # where X axis (recall) changes value
    i = np.where(mrec[:-1] != mrec[1:])[0]    

    # and sum (\Delta recall) * prec
    ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    return ap
def test_maps():
    with open('result_aux_out.pkl','rb') as f:
        data = pickle.load(f)
    frame = 195
    vs = [0.8,0.8]
    pcr = [-51.2, -51.2, -5.0, 51.2, 51.2, 3.0]
    score_threshold = 0.45
    rgb = ['r','y','b']
    for i in range(3):
        reg = data[frame]['maps'][i][0]['reg'].cpu().numpy()
        scores = data[frame]['maps'][i][0]['heatmap'].sigmoid().cpu().numpy()
        lines = get_lines_sets(reg,scores,pcr,vs,score_threshold,eps=0.2)
        print(len(lines))

def maps_eval():
    # 增加线的连续性 -> 同一条线断掉次数做评测
    import pickle
    with open('tmp/result/result_aux_470.pkl','rb') as f:
        data = pickle.load(f)
    gt_maps_all = []
    pred_maps_all = []
    pred_scores_all = []
    for d in data:
        gt_maps_all.append(d['gt_maps_3d'][0][0].cpu().numpy())
        pred_map, pred_sc = [], []
        for maps in d['maps']:
            pred_maps = maps['maps_points'].reshape(-1,2).cpu().numpy()
            pred_score = maps['maps_scores'].view(-1,1)
            pred_score = torch.cat([pred_score,pred_score],-1)
            pred_score = pred_score.cpu().numpy().reshape(-1)
            pred_map.append(pred_maps)
            pred_sc.append(pred_score)
        pred_maps_all.append(pred_map)
        pred_scores_all.append(pred_sc)
    mp, mr, map, mf1 = eval_map(gt_maps_all ,pred_maps_all, pred_scores_all,
             dist_threshold=[0.5, 1.0, 2.0],score_threhold=0.2)
    print(1)

def l1l2error(l1,l2,eps=1e-4):
    return abs(l1[0] - l2[0]) + abs(l1[1] - l2[1])

def divide_lines(l1, l2, eps=1e-4):
    l1 = np.array(l1)
    l2 = np.array(l2)
    c1 = (l1[:-1] + l1[1:])[:,:2]/2
    c2 = (l2[:-1] + l2[1:])[:,:2]/2
    c1 = c1.reshape(-1,1,2)
    c2 = c2.reshape(1,-1,2)
    cd = c1 - c2
    cd = np.abs(cd).sum(-1)
    coor1, coor2 = np.where(cd <= eps)
    lin = []
    flag = False
    if len(coor1) > 0:
        coor1 = coor1.tolist()
        coor2 = coor2.tolist()
        coor1 = [-1] + coor1 + [l1.shape[0]]
        coor2 = [-1] + coor2 + [l2.shape[0]]
        for i in range(len(coor1)-1):
            lin.append(l1[coor1[i]+1:coor1[i + 1]+1])
        
        for i in range(len(coor2)-1):
            lin.append(l2[coor2[i]+1:coor2[i + 1]+1])
        flag = True
    return lin,flag

def delete_replines(lines, eps=1e-4):
    len_lines = len(lines)
    if len_lines<2:
        return lines
    i = 0
    f = 0
    while i < len(lines)-1:
        j= i + 1
        while j < len(lines):
            l1, l2 = lines[i]['points'], lines[j]['points']
            new_line,flag = divide_lines(l1, l2)
            if flag:
                f = 1
                for new_l in new_line:
                    tmp = copy.deepcopy(lines[0])
                    if new_l.shape[0] > 1:
                        tmp['points'] = new_l.tolist()
                        lines.append(tmp)
                lines.pop(j)
                lines.pop(i)
            if f:
                break
            else:
                j += 1
        if f:
            i += 0
            f = 0
        else:
            i += 1
    return lines
def comb_lines(lines,eps=1e-4,reps=120):
    lin = []
    len_lines = len(lines)
    if len_lines < 2:
        return lines
    save_lines = [True for i in range(len_lines)]
    for i in range(len_lines-1):
        if not save_lines[i]:
            continue
        for j in range(i+1, len_lines):
            if (not save_lines[j]):
                continue
            # print(i,j,save_lines[i], save_lines[j])
            flag = 0
            if isinstance(lines[i], dict):
                l1, l2 = lines[i]['points'], lines[j]['points']
            else:
                l1, l2 = lines[i], lines[j]
            if l1l2error(l1[0],l2[0]) < eps:
                l1 = l1[::-1]
                flag = 1
            elif l1l2error(l1[0],l2[-1]) < eps:
                l1 = l1[::-1]
                l2 = l2[::-1]
                flag = 1
            elif l1l2error(l1[-1],l2[0]) < eps:
                flag = 1
            elif l1l2error(l1[-1],l2[-1]) < eps:
                l2 = l2[::-1]
                flag = 1
            else:
                flag = 0
            if flag==1:
                x11,y11 = l1[-2][:2]
                x12,y12 = l1[-1][:2]
                x21,y21 = l2[0][:2]
                x22,y22 = l2[1][:2]
                v1 = np.array([x12 - x11, y12 - y11])
                v2 = np.array([x22 - x21, y22 - y21])
                # 两个向量
                Lx=np.linalg.norm(v1,ord=2)
                Ly=np.linalg.norm(v2,ord=2)
                cos_angle=min(v1.dot(v2)/(Lx*Ly),1.0)
                r=np.arccos(cos_angle)
                if abs(r) >= reps * np.pi / 180:
                    continue
                l1.pop(-1)
                l1.extend(l2)
                if isinstance(lines[i], dict):
                    lines[i]['points'] = l1
                else:
                    lines[i] = l1
                # print(i,j)
                save_lines[j] = False
        lin.append(lines[i])
    return lin

def comb_lines_m(lines,eps=1e-4,reps=120):
    lin = []
    lines.sort(key=lambda x: len(x), reverse=True)
    len_lines = len(lines)
    if len_lines < 2:
        return lines
    save_lines = [True for i in range(len_lines)]
    for i in range(len_lines-1):
        if not save_lines[i]:
            continue
        for j in range(i+1, len_lines):
            if (not save_lines[j]):
                continue
            # print(i,j,save_lines[i], save_lines[j])
            flag = 0
            if isinstance(lines[i], dict):
                l1, l2 = lines[i]['points'], lines[j]['points']
            else:
                l1, l2 = lines[i], lines[j]
            mind = eps
            if l1l2error(l1[0],l2[0]) < mind:
                l1 = l1[::-1]
                flag = 1
                mind = l1l2error(l1[0],l2[0])
            elif l1l2error(l1[0],l2[-1]) < mind:
                l1 = l1[::-1]
                l2 = l2[::-1]
                flag = 1
                mind = l1l2error(l1[0],l2[-1])
            elif l1l2error(l1[-1],l2[0]) < mind:
                flag = 1
                mind = l1l2error(l1[-1],l2[0])
            elif l1l2error(l1[-1],l2[-1]) < mind:
                l2 = l2[::-1]
                flag = 1
                mind = l1l2error(l1[-1],l2[-1])
            else:
                flag = 0
            if flag==1:
                x11,y11 = l1[-2][:2]
                x12,y12 = l1[-1][:2]
                x21,y21 = l2[0][:2]
                x22,y22 = l2[1][:2]
                v1 = np.array([x12 - x11, y12 - y11])
                v2 = np.array([x22 - x21, y22 - y21])
                # 两个向量
                Lx=np.linalg.norm(v1,ord=2)
                Ly=np.linalg.norm(v2,ord=2)
                cos_angle=min(v1.dot(v2)/(Lx*Ly),1.0)
                r=np.arccos(cos_angle)
                if abs(r) > reps * np.pi / 180:
                    continue
                l1.pop(-1)
                l1.extend(l2)
                if isinstance(lines[i], dict):
                    lines[i]['points'] = l1
                else:
                    lines[i] = l1
                # print(i,j)
                save_lines[j] = False
        lin.append(lines[i])
    return lin
def decode_one_voxels(dt_pts):
    # c++ 换成if else 结构
    # dt_pts = dt_pts % 4
    dt_pts = np.clip(dt_pts,a_min=0, a_max=4)
    dx = (dt_pts <= 1) * 0 + \
            ((1 < dt_pts) & (dt_pts <= 2)) * (dt_pts - 1) + \
            ((2 < dt_pts) & (dt_pts <= 3)) * 1 + \
            (3 < dt_pts) * (4 - dt_pts)
    dy = (dt_pts <= 1) * dt_pts + \
            ((1 < dt_pts) & (dt_pts <= 2)) * 1 + \
            ((2 < dt_pts) & (dt_pts <= 3)) * (3 - dt_pts) + \
            (3 < dt_pts) * 0
    dir = (dt_pts <= 1) * 0 + \
            ((1 < dt_pts) & (dt_pts <= 2)) * 1 + \
            ((2 < dt_pts) & (dt_pts <= 3)) * 2 + \
            (3 < dt_pts) * 3
    return dx, dy, dir

def norm2(p1,p2):
    return ((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2) ** (0.5)

def decode_one_lines(i, j, r1, r2):
    vx = j
    vy = i
    d1x, d1y, dir1 = decode_one_voxels(0 + r1) 
    d2x, d2y, dir2 = decode_one_voxels(4 - r2)

    d1x = d1x + vx
    d1y = d1y + vy
    d2x = d2x + vx
    d2y = d2y + vy

    return [d1x,d1y,i,j,dir1],[d2x,d2y,i,j,dir2]

def serach_voxels(line, scores, reg, pv, 
                  score_threshold, mode=-1, eps=0.2):
    _,_,H,W = reg.shape
    p = line[mode]
    _,_,j,i,dir = p
    # next_dirs = [[[i-1,j],[i-1,j+1],[i-1,j-1],[i,j+1],[i,j-1]],
    #             [[i,j+1],[i-1,j+1],[i+1,j+1],[i-1,j],[i+1,j]],
    #             [[i+1,j],[i+1,j+1],[i+1,j-1],[i,j+1],[i,j-1]],
    #             [[i, j-1],[i-1,j-1],[i+1,j-1],[i-1,j],[i+1,j]]]
    next_dirs = [[[i-1,j],[i-1,j+1],[i-1,j-1]],
                [[i,j+1],[i-1,j+1],[i+1,j+1]],
                [[i+1,j],[i+1,j+1],[i+1,j-1]],
                [[i, j-1],[i-1,j-1],[i+1,j-1]]]
    flag = False
    mind,which_p,which_layer = 1e5, -1, -1 # error, 
    for next_dir in next_dirs[dir]:
        j_next, i_next = next_dir
        if (i_next < 0) | (j_next < 0) | (i_next >= W) | (j_next >= H):
            continue
        for k in range(pv):
            if scores[0,k,i_next,j_next] > score_threshold:
                r1 = reg[0, pv*k + 0, i_next, j_next]
                r2 = reg[0, pv*k + 1, i_next, j_next]
                p1, p2 = decode_one_lines(i_next,j_next,r1,r2)
                tmp_d = norm2(p1,p)
                if  tmp_d < mind:
                    mind = tmp_d
                    which_p = [p1, p2]
                tmp_d = norm2(p2,p)
                if  tmp_d < mind:
                    mind = tmp_d
                    which_p = [p2, p1]
                if mind < eps:
                    which_layer = k
    if which_layer > -1:
        _,_,i,j,_ = which_p[1]
        scores[0,which_layer,i,j] = 0
        line[mode][0] = (line[mode][0] + which_p[0][0])/2
        line[mode][1] = (line[mode][1] + which_p[0][1])/2
        if mode == -1:
            insert_id = len(line)
        else:
            insert_id = 0
        line.insert(insert_id, which_p[1])
    else:
        flag = True
    return line, scores, flag
    
def serach_voxels_roadedge(line, scores, reg, pv, 
                  score_threshold, mode=-1, eps=0.2,nms=True,reps=85,w=0.5):
    _,_,H,W = reg.shape

    p = line[mode] 
    if mode == -1:
        p_ = line[-2]
    else:
        p_ = line[1]
    _,_,j,i,dir = p[:6]
    # next_dirs = [[[i-1,j],[i-1,j+1],[i-1,j-1],[i,j+1],[i,j-1]],
    #             [[i,j+1],[i-1,j+1],[i+1,j+1],[i-1,j],[i+1,j]],
    #             [[i+1,j],[i+1,j+1],[i+1,j-1],[i,j+1],[i,j-1]],
    #             [[i, j-1],[i-1,j-1],[i+1,j-1],[i-1,j],[i+1,j]]]
    next_dirs = [[[i-1,j],[i-1,j+1],[i-1,j-1]],
                [[i,j+1],[i-1,j+1],[i+1,j+1]],
                [[i+1,j],[i+1,j+1],[i+1,j-1]],
                [[i, j-1],[i-1,j-1],[i+1,j-1]]]

    
    flag = False
    mind,which_p,which_layer,minr = 1e5, -1, -1, 1e5 # error, 
    # minds = []
    # which_ps = []
    for next_dir in next_dirs[dir]:
        j_next, i_next = next_dir
        if (i_next < 0) | (j_next < 0) | (i_next >= W) | (j_next >= H):
            continue
        for k in range(pv):
            if scores[0,k,i_next,j_next] > score_threshold:
                r1 = reg[0, pv*k + 0, i_next, j_next]
                r2 = reg[0, pv*k + 1, i_next, j_next]
                p1, p2 = decode_one_lines(i_next,j_next,r1,r2)
                tmp_d1 = norm2(p1,p)
                tmp_d2 = norm2(p2,p)

                if  tmp_d1 < tmp_d2:
                    mind1 = tmp_d1
                    tmp_which_p = [p1, p2]
                else:
                    mind1 = tmp_d2
                    tmp_which_p = [p2, p1]
                if mind1 < eps:
                    # minds.append(mind)
                    x11,y11 = tmp_which_p[0][:2]
                    x12,y12 = tmp_which_p[1][:2]
                    x21,y21 = p_[:2]
                    x22,y22 = p[:2]
                    v1 = np.array([x12 - x11, y12 - y11])
                    v2 = np.array([x22 - x21, y22 - y21])
                    # 两个向量
                    Lx=np.linalg.norm(v1,ord=2)
                    Ly=np.linalg.norm(v2,ord=2)
                    cos_angle=min(v1.dot(v2)/(Lx*Ly),1.0)
                    r=np.arccos(cos_angle)
                    if mind1 < mind and abs(r) < reps * np.pi / 180:
                        minr_ = mind1 / mind * w + (abs(r) * 180/np.pi/reps) * (1-w)
                        if minr_ < minr:
                            minr = minr_+0
                            which_layer = k
                            mind = mind1 + 0
                            which_p = copy.deepcopy(tmp_which_p)
                    # which_ps.append([p1, p2])
    # for l in which_ps: # tan(θ1) = abs((k2 - k1)/(1+k1*k2))
    #     k1 = l[1][-1] 
    #     tanr = abs((k1-k)/1 + k1*k)
    #     r = atan(tanr)
    #     if r < reps:
    #         pass
    if which_layer > -1:
        _,_,i,j,_ = which_p[1]
        scores[0,which_layer,i,j] = 0.0
        wr = 0.5
        line[mode][0] = (line[mode][0] * wr  + which_p[0][0] *(1-wr))
        line[mode][1] = (line[mode][1] * wr + which_p[0][1] * (1-wr))
        if mode == -1:
            insert_id = len(line)
        else:
            insert_id = 0
        # line.insert(insert_id, which_p[0])
        line.insert(insert_id, which_p[1])
        j,i,_dir = line[mode][2:5]
        nms_dirs = [[[i-1,j],[i-1,j+1],[i-1,j-1]],
                [[i,j+1],[i-1,j+1],[i+1,j+1]],
                [[i+1,j],[i+1,j+1],[i+1,j-1]],
                [[i, j-1],[i-1,j-1],[i+1,j-1]]]
        _nms_dir = nms_dirs[_dir]
        # _nms_dir = [[j,i] for i,j in _nms_dir]
        if nms:
            for next_dir in next_dirs[dir]:
                if next_dir  in _nms_dir:
                    continue
                j_next, i_next = next_dir
                if (i_next < 0) | (j_next < 0) | (i_next >= W) | (j_next >= H):
                    continue
                scores[0,:,i_next,j_next] = 0
    else:
        flag = True
    
    return line, scores, flag

def serach_voxels_2_dir(line, scores, 
                        reg, pv, score_threshold,
                        eps=0.2,nms=True):
    problem = 0
    while (not problem):
        line, scores, problem = serach_voxels_roadedge(line, scores, 
                reg, pv, score_threshold,
                mode=-1, eps=eps,nms=nms)
    problem = 0
    while (not problem):
        line, scores, problem = serach_voxels_roadedge(line, scores, 
                reg,  pv, score_threshold,
                mode=0, eps=eps,nms=nms)
    return line, scores
def get_lines_sets(reg, scores, score_threshold,score_m=0.1, eps=0.5, comb_eps=2.0,nms=True):
    _,pv,H,W = scores.shape
    lines_sets = []
    # for i in range(H/2):
    if 1:
        for i in range(0,H,1):
        # i = 30
            # if i ==30:
            for j in range(W):
                for k in range(pv):
                    line = []
                    if scores[0,k,i,j] > score_threshold:
                        # print(i,j,k)

                        r1 = reg[0, pv*k + 0, i, j]
                        r2 = reg[0, pv*k + 1, i, j]
                        p1, p2 = decode_one_lines(i,j,r1,r2)

                        line.append(p1)
                        line.append(p2)
                        
                        line, scores = serach_voxels_2_dir(line, scores, 
                                        reg, pv, score_threshold-score_m,
                                        eps=eps,nms=nms)
                        lines_sets.append(line)
    if 1:
        lines_sets = comb_lines_m(lines_sets,eps=comb_eps)
    return lines_sets

def updata_dir(dir,i,j):
    dir%=4
    if dir==0:
        i +=1
    if dir==1:
        j +=1
    if dir==2:
        i -=1
    if dir==3:
        j -=1
    return dir,i,j

def get_lines_sets_circle(reg, scores, score_threshold,score_m=0.1, eps=0.5,comb_eps=2.0):
    _,pv,H,W = scores.shape
    N = int(H/2)
    lines_sets = []
    i,j = N,N
    dir = 0
    for m in range(1,N*2+2):
        for s in range(2):
            for l in range(m):
                dir,i,j = updata_dir(dir,i,j)
                line = []
                if i >=H or j >=W:
                    continue
                for k in range(2):
                    if scores[0,k,i,j] > score_threshold:
                        # print(i,j,k)

                        r1 = reg[0, pv*k + 0, i, j]
                        r2 = reg[0, pv*k + 1, i, j]
                        p1, p2 = decode_one_lines(i,j,r1,r2)

                        line.append(p1)
                        line.append(p2)
                        
                        line, scores = serach_voxels_2_dir(line, scores, 
                                        reg, pv, score_threshold-score_m,
                                        eps=eps)
                        lines_sets.append(line)
            dir += 1
    if 1:
        lines_sets = comb_lines_m(lines_sets,eps=comb_eps)
    return lines_sets

def decode_lines_py(dt_cls, dt_pts,  ppv=2, scores_threshold=0.5):
    N, H, W, C = dt_cls.shape
    dt_pts = dt_pts.view(N,H,W,C,-1)
    xs = torch.arange(0, W, 1, dtype=torch.float32)
    ys = torch.arange(0, H, 1, dtype=torch.float32)

    xs, ys = torch.meshgrid([xs, ys])  # [x_grid, y_grid]
    xys = torch.stack((ys, xs), -1)

    # nxys = xys.unsqueeze(0).repeat(2,1,1,1)
    d0x, d0y, _ = decode_one_voxels(dt_pts[..., 0])
    d1y, d1x, _ = decode_one_voxels(dt_pts[..., 1])

    dxy = torch.stack([d0x, d0y, d1x, d1y], -1)
    dxy = dxy.reshape(N, H, W, C, ppv, 2) 
    xys = xys.reshape(1, H, W, 1, 1, 2)

    dxy = dxy + xys
    dxy = dxy.view(N, -1, ppv, 2)
    # dxy *=  torch.tensor(vs[:2]).cuda()
    
    dt_cls = dt_cls.view(N, -1)
    mask = dt_cls > scores_threshold
    scores = dt_cls[mask]
    lines_points = dxy[mask]
    return scores, lines_points
def test_process():
    import pickle
    import numpy as np
    file1 = "bevdetv2-nuscenes_infos_train_map.pkl"
    file0 = "bevdetv2-nuscenes_infos_val_map.pkl"
    root_path = '/media/gpu/sdc/datasets/nuscenes/'
    with open(root_path + file0,'rb') as f:
        data0 = pickle.load(f)
    # with open(root_path + file1,'rb') as f:
    #     data1 = pickle.load(f)
    frame=0
    # frame = np.random.randint(low=0,high=len(data0))

    data_frame = data0['infos'][frame]['vec_lines']['road_edge']
    # b = comb_lines(data_frame)

    b = delete_replines(copy.deepcopy(data_frame))

    print(len(b),len(data_frame))

    c = comb_lines(b)
    print(len(c),len(b))


def save_one(data, frame,show_file,save_dir):
    lines_type_rev = {"road_edge":0,
                    "dash_line":1,
                    "soild_line":2,
                    "cross_walk":3,
                    "center_line":4}
    lines_type = ["road_edge",
                  "solid_line",
                  "cross_walk",
                  "center_line"]
    lines_sets = []
    for i in range(4):
        scores = data[frame]['maps'][i][0]['heatmap'].sigmoid().cpu().numpy()
        scores = scores[:,:2]
        reg = data[frame]['maps'][i][0]['reg'].cpu().numpy()

        # reg = reg%4
        # scores[0][0] = morphology.skeletonize(scores[0][0]>0.4)
        # scores[0][1] = morphology.skeletonize(scores[0][1]>0.4)
        # lines_sets += get_lines_sets(reg, scores, 0.55, 0.1, 0.9, 4)
        nms = (i==0 or i == 3 or i == 0)
        lines_set = get_lines_sets(reg, scores, 0.55, 0.1, 0.9, 2, nms)
            # lines_sets
        for k, l in enumerate(lines_set):
            l = np.array(l)
            if len(l.shape) <2:
                continue
            if l.shape[0] < 4:
                continue
            l = l[:,:3]
            l *= 0.8
            l +=(-51.2)
            l[:,2] =0
            lines_sets+=[(l[:,:3],lines_type[i])]
    pts = data[frame]['points'][0].cpu().numpy()
    pts[:,3] *= 255
    pts_file = f"{frame:05d}"
    # save_dir = '00000'
    Toshow(points=pts,
        boxes=[], 
        lines=lines_sets,
        imgs=[], 
        calib=None,
        pcd_file=pts_file, 
        save_dir=save_dir,
        show_dir=show_file,
        is_bin=True)
def toann():
    from visualize import Toshow
    import os
    from tqdm import tqdm, trange

    for t in tqdm(range(100)):
        with open(f'../tmp/result/tmp/{t:06d}.pkl','rb') as f:
            data = pickle.load(f)
        for frame in range(len(data)):
            show_file = f'../tmp/toshow/'
            save_dir = f'{t:06d}'
            if not os.path.exists(show_file):
                os.mkdir(show_file)
            save_one(data, frame,show_file,save_dir)

def test_decode_lines():
    with open('./tmp/result/result_argov2_val_gap2000_all_decodeline.pkl','rb') as f:
        data = pickle.load(f)
    from skimage import morphology
    frame = 10
    A = []
    lines_sets = []
    for i in range(1):
        scores = data[frame]['maps'][i][0]['heatmap'].sigmoid().cpu().numpy()
        reg = data[frame]['maps'][i][0]['reg'].cpu().numpy()
        
        # lines_sets += get_lines_sets(reg, scores,score_threshold= 0.5,score_m= 0.1,eps=0.7)
        lines_sets += get_lines_sets(reg, scores, score_threshold= 0.5, score_m= 0.1, eps=0.7, comb_eps=2.0)
        # print(len(lines_sets[-1]))
    # for i in range(1):
    #     scores = data[frame]['maps'][i][0]['heatmap'].sigmoid().cpu().permute([0,2,3,1]).contiguous()
    #     reg = data[frame]['maps'][i][0]['reg'].cpu().permute([0,2,3,1]).contiguous()
        
    #     lines_sets +=decode_lines_py(scores, reg,  ppv=2, scores_threshold=0.5)
    #     print(len(lines_sets[-1]))
if __name__ == "__main__":
    # maps_eval()
    # test_process()
    test_decode_lines()